package com.codeboy.webshell.WebSocket.service.impl;

import com.codeboy.webshell.WebSocket.constant.ConstantPool;
import com.codeboy.webshell.WebSocket.pojo.SSHConnectInfo;
import com.codeboy.webshell.WebSocket.pojo.WebSSHData;
import com.codeboy.webshell.WebSocket.service.WebSSHservice;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.jcraft.jsch.Channel;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Arrays;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * @Description : WebSSH业务逻辑实现
 * @Author : 崔昌赫
 * @Date : 2021/9/12
 */
@Service
public class WebSSHserviceImpl implements WebSSHservice {
    private static Map<String,Object> sshMap = new ConcurrentHashMap<>(); //存放ssh连接信息的map
    private Logger logger = LoggerFactory.getLogger(WebSSHserviceImpl.class); //日志
    private ExecutorService executorService = Executors.newCachedThreadPool();//线程池

    /**
     * @Description : 初始化连接
     * @param session ssh会话实例
     * @Author : 崔昌赫
     * @Date : 2021/9/12
     */
    @Override
    public void initConnection(WebSocketSession session) {
        JSch jSch = new JSch();
        SSHConnectInfo sshConnectInfo = new SSHConnectInfo();
        sshConnectInfo.setjSch(jSch);
        sshConnectInfo.setWebSocketSession(session);
        String uuid = String.valueOf(session.getAttributes().get(ConstantPool.USER_UUID_KEY));
        sshMap.put(uuid,sshConnectInfo);
    }

    /**
     * @Description : 处理客户端发送的数据
     * @param buffer 接发送的信息数组
     * @param session ssh会话实例
     * @Author : 崔昌赫
     * @Date : 2021/9/12
     */
    @Override
    public void recvHandle(String buffer, WebSocketSession session) {
        ObjectMapper objectMapper=new ObjectMapper();
        WebSSHData webSSHData=null;
        try{
            webSSHData=objectMapper.readValue(buffer,WebSSHData.class);
        }catch (IOException e){
            logger.error("JSON转换异常!");
            logger.error("异常信息:{}",e.getMessage());
            return;
        }
        String userId=String.valueOf(session.getAttributes().get(ConstantPool.USER_UUID_KEY));
        if(ConstantPool.WEBSSH_OPERATE_CONNECT.equals(webSSHData.getOperate())){
            //查找刚刚存储的SSH实例
            SSHConnectInfo sshConnectInfo=(SSHConnectInfo) sshMap.get(userId);
            //启动线程异步处理
            WebSSHData finalWebSSHData =webSSHData;
            executorService.execute(new Runnable() {
                @Override
                public void run() {
                    try{
                        connectToSSH(sshConnectInfo,finalWebSSHData,session);
                    }catch (JSchException|IOException e){
                        logger.error("Webssh连接异常！异常信息:{}",e.getMessage());
                        close(session);
                    }
                }
            });
        }else if(ConstantPool.WEBSSH_OPERATE_COMMAND.equals(webSSHData.getOperate())){
            String command = webSSHData.getCommand();
            SSHConnectInfo sshConnectInfo=(SSHConnectInfo)sshMap.get(userId);
            if(sshConnectInfo!=null){
                try{
                    transToSSH(sshConnectInfo.getChannel(),command);
                }catch (IOException e){
                    logger.error("Webssh连接异常！异常信息:{}",e.getMessage());
                    close(session);
                }
            }

        }else{
            logger.error("不支持的操作");
            close(session);
        }

    }

    /**
     * @Description : 想客户端发送数据回显
     * @param session SSH会话实例
     * @param buffer 回写数据数组
     * @throws IOException
     * @Author : 崔昌赫
     * @Date : 2021/9/12
     */
    @Override
    public void sendMassage(WebSocketSession session, byte[] buffer) throws IOException {
        session.sendMessage(new TextMessage(buffer));

    }

    /**
     * @Description : 关闭WebSSH连接
     * @param session SSH会话实例
     * @Author : 崔昌赫
     * @Date : 2021/9/12
     */
    @Override
    public void close(WebSocketSession session) {
        String userId = String.valueOf(session.getAttributes().get(ConstantPool.USER_UUID_KEY));
        SSHConnectInfo sshConnectInfo=(SSHConnectInfo) sshMap.get(userId);
        if(sshConnectInfo!=null){
            //断开连接
            if(sshConnectInfo.getChannel()!=null) sshConnectInfo.getChannel().disconnect();
            //从map中移除
            sshMap.remove(userId);
        }

    }

    /**
     * @Description : 建立websocket并且连接到服务器ssh
     * @param sshConnectInfo SSH连接状态存储类
     * @param webSSHData SSH连接数据存储类
     * @param webSocketSession WebSocket连接实例
     * @throws JSchException 与SSH服务器连接失败抛出错误
     * @throws IOException IO错误
     * @Author : 崔昌赫
     * @Date : 2021/9/12
     */

    private void connectToSSH(SSHConnectInfo sshConnectInfo,WebSSHData webSSHData,WebSocketSession webSocketSession) throws JSchException,IOException{
        Session session=null;
        Properties config = new Properties();
        config.put("StrictHostKeyChecking","no");
        //获取jsch的会话
        session=sshConnectInfo.getjSch().getSession(webSSHData.getUsername(),webSSHData.getHost(),webSSHData.getPort());
        session.setConfig(config);
        //设置密码
        session.setPassword(webSSHData.getPassword());
        //设置连接超时时间30s
        session.connect(30000);
        //开启shell通道
        Channel channel=session.openChannel("shell");
        //设置通道连接超时时间3s
        channel.connect(3000);
        //设置channel
        sshConnectInfo.setChannel(channel);
        //转发消息
        transToSSH(channel,"clear\r");
        //读取终端返回的信息流
        InputStream inputStream = channel.getInputStream();
        try{
            //循环读取
            byte[] buffer=new byte[1024];
            int i =0;
            //如果没有数据来，线程会一直阻塞在这个地方等待数据。
            while((i=inputStream.read(buffer))!=-1){
                sendMassage(webSocketSession, Arrays.copyOfRange(buffer,0,i));
            }
        }finally {
            //断开连接后关闭会话
            session.disconnect();
            channel.disconnect();
            if(inputStream!=null){
                inputStream.close();
            }
        }


    }

    /**
     * @Description : 从后端向服务器SSH发送命令
     * @param channel WebSSH通道实例
     * @param command 发送的命令
     * @throws IOException IO错误
     * @Author : 崔昌赫
     * @Date : 2021/9/12
     */
    private void transToSSH(Channel channel,String command)throws IOException{
        if(channel!=null){
            OutputStream outputStream=channel.getOutputStream();
            outputStream.write(command.getBytes());
            outputStream.flush();
        }
    }

}
